package com.Q04;

import cn.hutool.core.lang.Assert;
import com.common.Prie;
import com.common.TreeNode;

import java.util.HashMap;

/**
 * @version v1.0.0
 * 创建时间：2022/7/8 10:13
 */
public class Q04_12 {
    public int pathSum(TreeNode root, int sum) {
        return pathSumPrie(root,sum);
    }
    private int pathSumPrie(TreeNode root, int sum){
        // key 为距离，value 为 从root开始的 key走法数量。例如2,5 从root走2距离，有5种走法
        HashMap<Long, Integer> prefix = new HashMap<Long, Integer>();
        prefix.put(0L, 1);
        return dfs(root, prefix, 0, sum);
    }

    private int dfs(TreeNode root,HashMap<Long, Integer> map,long cur,int sum){
        if (root==null) return 0;
        cur += root.val;
        int res = map.getOrDefault(cur-sum,0);
        map.put(cur,map.getOrDefault(cur,0)+1);
        res+=dfs(root.left,map,cur,sum);
        res+=dfs(root.right,map,cur,sum);
        map.put(cur,map.getOrDefault(cur,0)-1);
        return res;
    }

    private int pathSumRecur(TreeNode root, int sum){
        if(root==null) return 0;
        int ret = pathRecur(root, sum);
        return pathSumRecur(root.left,sum)+pathSumRecur(root.right,sum)+ret;
    }

    private int pathRecur(TreeNode root, int sum){
        int ret =0;
        if(root==null) return 0;
        if(root.val==sum)
            ret++;
        int left =pathRecur(root.left,sum-root.val);
        int right = pathRecur(root.right,sum-root.val);
        ret = ret+left+right;
        return ret;
    }

    public static void main(String[] args) {
        TreeNode treeNode = TreeNode.buildTree(new Integer[]{5, 4, 8, 11, null, 13, 4, 7, 2, null, null, 5, 1});
        Assert.isTrue(new Q04_12().pathSum(treeNode, 22)==3);
        Assert.isTrue(new Q04_12().pathSum(new TreeNode(0), 0)==1);
        Assert.isTrue(new Q04_12().pathSum(new TreeNode(1), 0)==0);
    }
}
